
import sys
from openft.open_quant_context import *
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
import codecs
import json

RET_OK = 0
class StockQuoteTest(StockQuoteHandlerBase):
    def on_recv_rsp(self, rsp_str):
        ret_code, content = super(StockQuoteTest, self).on_recv_rsp(rsp_str)
        if ret_code != RET_OK:
            print("StockQuoteTest: error, msg: %s" % content)
            return RET_ERROR, content
        print("StockQuoteTest ", content)
        return RET_OK, content

class OrderBookTest(OrderBookHandlerBase):
    def on_recv_rsp(self, rsp_str):
        ret_code, content = super(OrderBookTest, self).on_recv_rsp(rsp_str)
        if ret_code != RET_OK:
            print("OrderBookTest: error, msg: %s" % content)
            return RET_ERROR, content
        print("OrderBookTest", content)
        return RET_OK, content

class CurKlineTest(CurKlineHandlerBase):
    def on_recv_rsp(self, rsp_str):
        ret_code, content = super(CurKlineTest, self).on_recv_rsp(rsp_str)
        if ret_code != RET_OK:
            print("CurKlineTest: error, msg: %s" % content)
            return RET_ERROR, content
        print("CurKlineTest", content)
        return RET_OK, content

class TickerTest(TickerHandlerBase):
    def on_recv_rsp(self, rsp_str):
        ret_code, content = super(TickerTest, self).on_recv_rsp(rsp_str)
        if ret_code != RET_OK:
            print("TickerTest: error, msg: %s" % content)
            return RET_ERROR, content
        print("TickerTest", content)
        return RET_OK, content

#连接服务器
def _dbhelper():
    quote_context = OpenQuoteContext(host='127.0.0.1', async_port=11111)
    quote_context.set_handler(StockQuoteTest())
    quote_context.set_handler(OrderBookTest())
    quote_context.set_handler(CurKlineTest())
    quote_context.set_handler(TickerTest())
    return quote_context

# 取符合市场和股票类型条件的股票简要信息
# market: 市场标识, string，例如，”HK”，”US”；具体见市场标识说明
# stock_type: 证券类型, string, 例如，”STOCK”，”ETF”；具体见证券类型说明
def _example_stock_basic(market,stock_type):

        quote_ctx =_dbhelper()
        ret_status, ret_data = quote_ctx.get_stock_basicinfo(market, stock_type)
        if ret_status == RET_ERROR:
            print('服务器繁忙！')
            exit()
        else:
            ret1 = ret_data[0:2]
            base = []
            result = []
            base = ret_Json_base(ret1)
            result = base.rstrip(',') +"]"
            return  result

#股票信息拼接Json
def ret_Json_base(ret1):
    bas ="["
    rets = []
    for item in ret1:
        bas += '{'
        bas += '"code":"' + item["code"]+'",'
        bas += '"lot_size":"' + str(item["lot_size"]) + '",'
        bas += '"stock_type":"' + item["stock_type"] + '",'
        bas += '"name":"' + item["name"] + '",'
        bas +=   ret_Json_snapic(_get_market_snapshot(str_list(item["code"])))
        rets.append(bas)

    return bas

#市场快照拼接Json
def ret_Json_snapic(ret2):
    bas = ""
    for item2 in ret2:
        # print(ss["code"])
        bas += '"prev_close_price":"' + str(item2["prev_close_price"])+ '",'
        bas += '"suspension":"'+ str(item2["suspension"]) + '",'
        bas += '"low_price":"'+ str(item2["low_price"]) + '",'
        bas += '"volume":"'+ str(item2["volume"]) + '",'
        bas += '"turnover_rate":"' + str(item2["turnover_rate"]) + '",'
        bas += '"high_price":"'+ str(item2["high_price"]) + '",'
        bas += '"last_price":"'+ str(item2["last_price"]) + '",'
        bas += '"turnover":"' + str(item2["turnover"]) + '",'
        bas += '"data_date":"' + str(item2["data_date"]) + '",'
        bas += '"data_time":"' + str(item2["data_time"]) + '",'
        bas += '"open_price":"' + str(item2["open_price"]) + '"}'
        bas += ","


    return  bas

#实时K线拼接Json
def ret_Json_Cur(ret2):
    bas = ""
    for item2 in ret2:
        bas += '"volume":"' + str(item2["volume"])+ '",'
        bas += '"turnover":"'+ str(item2["suspension"]) + '",'
        bas += '"low":"'+ str(item2["low"]) + '",'
        bas += '"k_type":"'+ str(item2["k_type"]) + '",'
        bas += '"close":"' + str(item2["close"]) + '",'
        bas += '"high":"'+ str(item2["high_price"]) + '",'
        bas += '"open":"'+ str(item2["open"]) + '",'
        bas += '"time_key":"' + str(item2["time_key"]) + '"}'
        bas += ","
    return  bas
#股票信息（含快照，k线图数据）
def base_snapic_Cur(ret1):
    bas = "["
    rets = []
    for item in ret1:
        bas += '{'
        bas += '"code":"' + item["code"] + '",'
        bas += '"lot_size":"' + str(item["lot_size"]) + '",'
        bas += '"stock_type":"' + item["stock_type"] + '",'
        bas += '"name":"' + item["name"] + '",'
        bas += ret_Json_snapic(_get_market_snapshot(str_list(item["code"])))
        bas += ret_Json_Cur(get_cur_return(str_list(item["code"]),item["stock_type"],1))

        rets.append(bas)
#获取K线图数据
def get_cur_return(code,data_type,num):
    db = _dbhelper()
    Subscribe(code, data_type)
    ret_code, ret_date = db.get_cur_kline(code, num)
    return ret_date

 # 获取市场快照 get_market_snapshot
#获取市场快照 get_market_snapshot有返回值
def _get_market_snapshot(codelist):
    #code_list = ["HK.00700"]
    quote_ctx = _dbhelper()
    ret_code, ret_data = quote_ctx.get_market_snapshot(codelist)
    return ret_data
#获取k线图
def _example_cur_kline(code_list,types):

    stock_code_list = [code_list]
    sub_type_list = [types]
    quote_ctx = _dbhelper()
    for code in stock_code_list:
        for sub_type in sub_type_list[0]:
            ret_status, ret_data = quote_ctx.subscribe(code, sub_type)
            if ret_status != RET_OK:
                print("%s %s: %s" % (code, sub_type, ret_data))
                exit()
    ret_status, ret_data = quote_ctx.query_subscription()
    if ret_status == RET_ERROR:
        print(ret_data)
        exit()

    for code in stock_code_list:
        for ktype in types:
            ret_code, ret_data = quote_ctx.get_cur_kline(code,1000, ktype)
            if ret_code == RET_ERROR:
                print(code, ktype, ret_data)
                exit()
            kline_table = ret_data
            #print("%s KLINE %s" % (code, ktype))
            #print(kline_table)
            #print("\n\n")

    return  ret_data

#字符串转list
def str_list(str):
     files = str.replace('[', '').replace(']', '')
     files = files.split(',')
     files_List = []
     for i in range(len(files)):
         files[i] = files[i].replace("'", '')
         files_List.append(files[i])

     return  files_List
#订阅
def Subscribe(code,data_type):
    db = _dbhelper()
    ret_code, ret_data = db.subscribe(code, data_type)
    if ret_code != RET_OK:
        print("订阅失败！")
        exit()
    else:
        return ret_data
#退订
def Unsubscribe(code,data_type):
    db = _dbhelper()
    ret_code, ret_data = db.unsubscribe(code, data_type)
    if ret_code != RET_OK:
        print("退订失败！")
        exit()
    else:
        return ret_data

def _get_cur_kline_basicInfo(code,num,stock_type):
    quote_ctx = _dbhelper()
    Subscribe(code,stock_type)
    ret_status, ret_data = quote_ctx.get_cur_kline(code,num,stock_type)
    if ret_status == RET_ERROR:
        print('')
        exit()
    else:
        #ret1 = ret_data[0:1]
        #print(ret_data)
        #base = []
        #result = []
        # = base_snapic_Cur(ret1)
        #result = base.rstrip(',') + "]"
        return  ret_data

if __name__ == "__main__":

    #action = sys.argv[1]
    action = "ckBasic"
    if action == "basicInfo":
        market = sys.argv[2]
        stock_type = sys.argv[3]
        #获取股票列表
        res = _example_stock_basic(market,stock_type)
        print(res)


    elif action == "kline":
        _market = sys.argv[2]
        _types = sys.argv[3]
        num = sys.argv[4]
        autype = sys.argv[5]
        files_List = str_list(_types)
        #获取股票K线图
        res  = _example_cur_kline(_market,files_List)

    elif action == "snapshot":
        _market = sys.argv[2]
        if _market != "":
			 #获取市场快照
             res = _get_market_snapshot(str_list(_market))
             print(res)
        else:
            print("参数不能为空！")
    elif action =="ckBasic":
        #_market = sys.argv[2]
        #_num = sys.argv[3]
        # 获取实时K线
        mark = "HK.08300"
        res = _get_cur_kline_basicInfo(mark,1,"K_DAY")
        print(res)